import os
import sys
import shutil
import subprocess
import glob
import re
import time
import platform
import typer
import lief
from typing import Tuple, List, Optional


CORE_DUMP_PATH = "/tmp/dumped"
CORE_DUMP_PATH_TMP = f"{CORE_DUMP_PATH}.tmp"


def check_macos_version() -> None:
    typer.echo("[+] Checking macOS version...")
    if platform.system() != "Darwin":
        typer.echo("[-] This script only works on macOS")
        sys.exit(1)

    version = platform.mac_ver()[0]
    version_parts = list(map(int, version.split(".")))

    if version_parts[0] != 15 or version_parts[1] < 0 or version_parts[1] > 2:
        typer.echo(f"[-] This script requires macOS 15.0-15.2 (current version: {version})")
        sys.exit(1)

    typer.echo(f"[+] macOS version {version} confirmed")


def check_tools() -> None:
    typer.echo("[+] Checking required tools...")

    if subprocess.run(["which", "rabin2"], stdout=subprocess.PIPE).returncode != 0:
        typer.echo("[-] radare2 is not installed")
        typer.echo("    Please install with: brew install radare2")
        sys.exit(1)

    typer.echo("[+] All required tools are available")


def check_root() -> None:
    if os.geteuid() != 0:
        typer.echo("This script must be run as root")
        sys.exit(1)


def cleanup() -> None:
    typer.echo("[+] Cleaning up...")
    for file in [CORE_DUMP_PATH, CORE_DUMP_PATH_TMP]:
        if os.path.exists(file):
            os.remove(file)


def get_nonwritable_regions(pid: int) -> List[Tuple[str, int, int, str]]:
    vmmap_output = subprocess.run(["vmmap", "--wide", str(pid)], stdout=subprocess.PIPE, stderr=subprocess.PIPE).stdout.decode()
    pattern = re.compile(
        r'^(\S.*?\S)\s+([0-9a-fA-F]+)-([0-9a-fA-F]+)\s+\[.*?\]\s+\S+/\S+\s+\S+(?:\s+(.*))?$'
    )

    nonwritable_regions = []
    nonwritable_region_start = False
    for line in vmmap_output.split("\n"):
        if line.startswith("==== Non-writable regions for process"):
            nonwritable_region_start = True
            continue

        if line.startswith("==== Writable regions for process"):
            nonwritable_region_start = False
            continue

        if nonwritable_region_start:
            match = pattern.match(line)
            if match:
                region_type, start, end, detail = match.groups()
                nonwritable_regions.append((region_type, int(start, 16), int(end, 16), (detail or "").strip()))
    return nonwritable_regions


def launch_process(target_process: str) -> int:
    typer.echo("[+] Launching target process...")
    subprocess.run(["open", "-j", "-g", "-a", target_process])
    time.sleep(5)
    process = subprocess.Popen(["pgrep", "-x", target_process], stdout=subprocess.PIPE)
    pid = int(process.stdout.read().decode().strip())
    return pid


def dump_core(pid: int) -> str:
    typer.echo("[+] Dumping core...")
    subprocess.run(["gcore", "-d", "-v", "-o", CORE_DUMP_PATH, "-s", str(pid)])

    if os.path.exists(CORE_DUMP_PATH) and os.path.getsize(CORE_DUMP_PATH) > 0:
        typer.echo(f"[+] Core file saved to {CORE_DUMP_PATH}")
        return CORE_DUMP_PATH

    if os.path.exists(CORE_DUMP_PATH_TMP) and os.path.getsize(CORE_DUMP_PATH_TMP) > 0:
        typer.echo(f"[+] Core file saved to {CORE_DUMP_PATH_TMP}")
        return CORE_DUMP_PATH_TMP

    typer.echo("[-] Failed to dump core file")
    sys.exit(1)


def is_encrypted_binary(binary_path: str) -> Tuple[bool, lief.Binary]:
    try:
        binary = lief.parse(binary_path)
        if binary is None:
            return False, None
        if binary.format != lief.Binary.FORMATS.MACHO:
            return False, None
        for cmd in binary.commands:
            if cmd.command == lief.MachO.LoadCommand.TYPE.ENCRYPTION_INFO_64 and cmd.crypt_id == 1:
                return True, binary
        return False, None
    except Exception as e:
        typer.echo(f"[-] Failed to parse binary: {e}")
        return False, None


def get_app_path(target_process: str) -> str:
    typer.echo("[+] Getting app bundle path...")
    app_path = os.path.join("/Applications", f"{target_process}.app")
    if not os.path.exists(app_path):
        typer.echo("[-] Application not found")
        sys.exit(1)

    typer.echo(f"[+] Application path: {app_path}")
    if not os.path.exists(os.path.join(app_path, "Wrapper")):
        typer.echo("[-] Application is not an iOS app")
        sys.exit(1)
    
    return app_path


def get_encrypted_binaries(app_path: str) -> List[Tuple[str, lief.Binary]]:
    typer.echo("[+] Finding encrypted binaries...")
    encrypted_binaries = []
    for root, _, files in os.walk(app_path):
        for file in files:
            app_path = os.path.join(root, file)
            is_encrypted, binary = is_encrypted_binary(app_path)
            if is_encrypted:
                encrypted_binaries.append((app_path, binary))

    if len(encrypted_binaries) == 0:
        typer.echo("[-] No encrypted binaries found")
        sys.exit(1)

    return encrypted_binaries


def get_crypto_info(binary: lief.Binary) -> Tuple[int, int]:
    typer.echo("[+] Getting cryptoff and cryptsize...")

    for cmd in binary.commands:
        if cmd.command == lief.MachO.LoadCommand.TYPE.ENCRYPTION_INFO_64:
            typer.echo(f"[+] Cryptoff: {cmd.crypt_offset}")
            typer.echo(f"[+] Cryptsize: {cmd.crypt_size}")
            return cmd.crypt_offset, cmd.crypt_size

    typer.echo("[-] Failed to get cryptoff and cryptsize")
    sys.exit(1)


def to_ios_path(path: str) -> Optional[str]:
    wrapper_index = path.find("/Wrapper")
    if wrapper_index != -1:
        return path[wrapper_index + len("/Wrapper"):]
    return None


def get_decrypted_content(encrypted_binary_path: str, encrypted_binary: lief.Binary, core_file: lief.Binary, nonwritable_regions: List[Tuple[str, int, int, str]], cryptoff: int, cryptsize: int) -> Optional[bytes]:
    typer.echo("[+] Searching for decrypted content...")

    ios_path = to_ios_path(encrypted_binary_path)
    if ios_path is None:
        typer.echo("[-] Failed to convert path to iOS path")
        sys.exit(1)

    segment_content = b""
    segment_with_encrypted_content = encrypted_binary.segment_from_offset(cryptoff)
    typer.echo(f"[+] Segment with encrypted content: {segment_with_encrypted_content}")
    for region_type, start, _, detail in nonwritable_regions:
        if region_type == segment_with_encrypted_content.name and detail.endswith(ios_path):
            typer.echo(f"[+] Found decrypted content in {region_type} at {start}")
            segment = core_file.segment_from_virtual_address(start)
            typer.echo(f"[+] Segment: {segment}")
            segment_content += bytearray(segment.content)

    if segment_content:
        return segment_content[cryptoff - segment_with_encrypted_content.file_offset:cryptoff + cryptsize - segment_with_encrypted_content.file_offset]
    return None


def get_cryptid_offset(executable_path):
    typer.echo("[+] Getting cryptid offset...")
    process = subprocess.Popen(
        ["rabin2", "-I", executable_path], stdout=subprocess.PIPE
    )
    baddr = int(
        re.search(r"baddr\s+0x([0-9a-f]+)", process.stdout.read().decode()).group(1), 16
    )
    typer.echo(f"[+] Baddr: {baddr}")

    process = subprocess.Popen(
        ["rabin2", "-H", executable_path], stdout=subprocess.PIPE
    )
    cryptid_vaddr = int(
        re.search(r"0x([0-9a-f]+)\s+cryptid\s+", process.stdout.read().decode()).group(
            1
        ),
        16,
    )
    typer.echo(f"[+] Cryptid vaddr: {cryptid_vaddr}")

    return cryptid_vaddr - baddr


def create_ipa_file(
    app_path: str,
    decrypted_binaries: List[Tuple[str, str]],
    signing_identity: str,
    profile_path: str,
    entitlements_path: str,
) -> None:
    app_exec_name = os.path.basename(app_path).split(".")[0]
    dest_app_path = f"./Payload/{os.path.basename(app_path)}"

    typer.echo("[+] Removing previous payload directory...")
    subprocess.run(["rm", "-rf", "Payload"], check=True)
    subprocess.run(["rm", "-f", f"{app_exec_name}.ipa"], check=True)

    login_user = os.getlogin()

    typer.echo("[+] Copying original app content...")
    subprocess.run(["sudo", "-u", login_user, "mkdir", "Payload"], check=True)
    subprocess.run(["sudo", "-u", login_user, "cp", "-R", app_path, dest_app_path], check=True)

    typer.echo("[+] Copying decrypted executables...")
    for ios_path, decrypted_executable_path in decrypted_binaries:
        typer.echo(f"[+] Copying {decrypted_executable_path} to {os.path.join('./Payload', ios_path)}")
        subprocess.run(
            ["sudo", "-u", login_user, "cp", decrypted_executable_path, os.path.join("./Payload", ios_path)],
            check=True
        )

    typer.echo("[+] Copying profile...")
    subprocess.run(
        ["sudo", "-u", login_user, "cp", profile_path, os.path.join(dest_app_path, "embedded.provisionprofile")],
        check=True
    )

    typer.echo("[+] Re-signing frameworks...")
    for framework_path in glob.glob(os.path.join(dest_app_path, "Frameworks", "*")):
        subprocess.run(
            ["sudo", "-u", login_user, "codesign", "--force", "--sign", signing_identity, "--entitlements", entitlements_path, framework_path],
            check=True
        )

    typer.echo("[+] Re-signing main executable...")
    subprocess.run(
        ["sudo", "-u", login_user, "codesign", "--force", "--sign", signing_identity, "--entitlements", entitlements_path, dest_app_path],
        check=True
    )

    typer.echo("[+] Creating ipa file...")
    subprocess.run(
        ["sudo", "-u", login_user, "zip", "-r", f"{app_exec_name}.ipa", "Payload"],
        check=True
    )
    typer.echo(f"[+] IPA file created: {app_exec_name}.ipa")


def main(
    target_process: str = typer.Argument(..., help="Target process name"),
    signing_identity: str = typer.Option(
        ..., "--signing-identity", "-s", help="Signing identity to use for re-signing"
    ),
    profile_path: str = typer.Option(
        ..., "--profile", "-p", help="Path to provisioning profile"
    ),
    entitlements_path: str = typer.Option(
        ..., "--entitlements", "-e", help="Path to entitlements file"
    ),
):
    """
    Decrypt FairPlay-protected iOS apps and create a re-signed IPA file.
    """
    typer.echo("[+] Checking requirements...")
    check_root()
    check_macos_version()
    check_tools()
    cleanup()

    pid = launch_process(target_process)
    app_path = get_app_path(target_process)
    nonwritable_regions = get_nonwritable_regions(pid)

    core_file_path = dump_core(pid)
    core_file = lief.parse(core_file_path)

    encrypted_binaries = get_encrypted_binaries(app_path)
    decrypted_binaries = []
    for binary_path, binary in encrypted_binaries:
        typer.echo(f"[+] Processing {binary_path}")
        cryptoff, cryptsize = get_crypto_info(binary)
        decrypted_content = get_decrypted_content(binary_path, binary, core_file, nonwritable_regions, cryptoff, cryptsize)
        if decrypted_content is None:
            typer.echo("[-] Failed to get decrypted content (the binary does not appear to be contained in the core file)")
            continue
        typer.echo("[+] Decrypted content found")

        output_path = f"/tmp/{os.path.basename(binary_path)}.decrypted"
        typer.echo(f"[+] Copying original binary {binary_path} to {output_path}")
        shutil.copy(binary_path, output_path)

        typer.echo("[+] Copying decrypted content...")
        with open(output_path, "r+b") as dst:
            dst.seek(cryptoff)
            dst.write(decrypted_content)

        cryptid_offset = get_cryptid_offset(binary_path)
        typer.echo(f"[+] Cryptid offset: {cryptid_offset}")
        typer.echo("[+] Patching cryptid...")
        with open(output_path, "r+b") as f:
            f.seek(cryptid_offset)
            f.write(b"\x00")
        typer.echo(f"[+] Decrypted executable saved to {output_path}")
        decrypted_binaries.append((to_ios_path(binary_path)[1:], output_path))

    typer.echo("[+] Creating ipa file...")
    wrapped_app_path = os.path.join(app_path, "Wrapper", os.path.basename(app_path))
    create_ipa_file(
        wrapped_app_path, decrypted_binaries, signing_identity, profile_path, entitlements_path
    )


if __name__ == "__main__":
    typer.run(main)
